from diffusers import DiffusionPipeline
from datetime import datetime
import torch
import gc


class Text2ImageModel:
  def __init__(self, model_id: str):
    pipeline = DiffusionPipeline.from_pretrained(model_id, use_safetensors=True)
    pipeline = pipeline.to("cuda")
    self.pipeline = pipeline
  
  def generate(self, prompt: str):
    image = self.pipeline(prompt).images[0] 
    return image
  
  def save_image(self, prompt: str, output_dir: str = ""):
    image = self.generate(prompt)
    ts = datetime.now().strftime("%Y%m%d%H%M%S")
    filename = f"{output_dir}{ts}.jpg"
    image.save(filename)
  
  def close(self):
    del self.pipeline
    gc.collect()
    torch.cuda.empty_cache()

__Text2ImageModelInstance__: Text2ImageModel = None

def text2image(model_id: str, prompt: str):
  global __Text2ImageModelInstance__
  if __Text2ImageModelInstance__ is None:
    print("create module instance..")
    __Text2ImageModelInstance__ = Text2ImageModel(model_id)
  return __Text2ImageModelInstance__.generate(prompt)

def text2image_once(model_id: str, prompt: str):
  model = Text2ImageModel(model_id)  
  image = model.generate(prompt)
  model.close()
  print("create module instance once..")
  return image

def text2imageN(model_id: str, prompt: str, n: int):
  model = Text2ImageModel(model_id)  
  # image = model.generate(prompt)
  images = []
  for _ in range(n):
    images.append(model.generate(prompt))
  model.close()
  print(f"create module instance {n} times..")
  return images